from stable_baselines3.common.env_checker import check_env
from my_env import CustomEnv
import argparse
import time
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import unwrap_wrapper, is_wrapped
from stable_baselines3.common.vec_env import DummyVecEnv
import numpy as np
import argparse
import cv2 as cv
import rospy
import math
import time
import rosgraph
from std_msgs.msg import Header
from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError



class RKNN_Agent():
    def __init__(self, args):
        # args settings
        self.tty_usb = args.tty_usb
        self.max_steps = args.max_steps
        self.rollout_steps = args.rollout_steps
        self.batch_size = args.batch_size
        self.log_path = args.log_path
        self.save_path = args.save_path
        self.model = args.model

        # ros settings
        self._bridge = CvBridge()
        self._subscriber = rospy.Subscriber('/usb_cam/image_raw',Image,self.callback,queue_size=1,buff_size=1000000)

        # reward setttings
        self._discount_factor = 7.0
        self._frame_count = 0

        # model settings
        
        self._env = CustomEnv(self.tty_usb, self._discount_factor)
        unwrap_wrapper(self._env, DummyVecEnv)
        

        # self._model = PPO.load(self.model,
        #                        env=self._env,
        #                        tensorboard_log=self.log_path,
        #                        )
        self._model = PPO('CnnPolicy',env=self._env,verbose=1)


    def callback(self,img_data):
        self._frame_count += 1

        # Image resize
        cv_img = self._bridge.imgmsg_to_cv2(img_data,"passthrough")
        gray_img = cv.cvtColor(cv_img, cv.COLOR_BGR2GRAY)
        show_img = cv.resize(gray_img,(320,240))
        cv.imshow("origin_img",show_img)
        cv.waitKey(1)

        self._env.forward_img(gray_img)


if __name__ == '__main__':
    parser = argparse.ArgumentParser('arg parser')
    parser.add_argument('--max_steps', type=int, default=25600)
    parser.add_argument('--batch_size', type=int, default=64)
    parser.add_argument('--rollout_steps', type=int, default=256)
    parser.add_argument('--model', type=str, default='model.zip')
    parser.add_argument('--save_path', type=str, default='./model')
    parser.add_argument('--log_path', type=str, default='./log')
    parser.add_argument('--tty_usb', type=str, default='/dev/ttyUSB2')
    args = parser.parse_args()

    rospy.init_node('rknn_control', log_level=rospy.ERROR)
    local_time = time.strftime("%m-%d-%H-%M")
    args.save_path = args.save_path + '/' + local_time
    args.log_path = args.log_path + '/' + local_time

    agent = RKNN_Agent(args)
    agent._model.learn(total_timesteps=10000,
                        )
    agent._model.save(agent.save_path)

    rospy.spin()

    agent.close()
